import numpy as np
import matplotlib.pyplot as plt
# 中文显示设置
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei", "SimSun"]
plt.rcParams["axes.unicode_minus"] = False
from data_handler import get_multi_asset_data
from trading_env import TradingEnv
from dqn_agent import DQNAgent
from config import WINDOW_SIZE, ASSET_CODES, TARGET_UPDATE


def calculate_metrics(asset_history):
    """计算回测指标：年化收益率、夏普比率、最大回撤"""
    total_assets = [h["total_asset"] for h in asset_history]
    returns = np.diff(total_assets) / total_assets[:-1]  # 日收益率

    # 年化收益率（假设每年252个交易日）
    annual_return = (total_assets[-1] / total_assets[0]) ** (252 / len(total_assets)) - 1

    # 夏普比率（无风险利率假设为3%）
    risk_free_rate = 0.03
    sharpe_ratio = (np.mean(returns) - risk_free_rate / 252) / np.std(returns) * np.sqrt(252)

    # 最大回撤
    cumulative_max = np.maximum.accumulate(total_assets)
    drawdown = (total_assets - cumulative_max) / cumulative_max
    max_drawdown = np.min(drawdown)

    return {
        "年化收益率": round(annual_return * 100, 2),
        "夏普比率": round(sharpe_ratio, 2),
        "最大回撤": round(max_drawdown * 100, 2)
    }


def backtest_buy_and_hold(data):
    """回测基线策略：买入并持有"""
    initial_cash = 100000
    n_assets = len(ASSET_CODES)
    equal_weight = 1 / n_assets  # 等权重配置
    buy_amount = (initial_cash * equal_weight) // data.iloc[WINDOW_SIZE][f"close_{list(ASSET_CODES.keys())[0]}"]

    positions = np.ones(n_assets) * buy_amount
    cash = initial_cash - np.sum(
        positions * [data.iloc[WINDOW_SIZE][f"close_{asset_type}"] for asset_type in ASSET_CODES.keys()])

    history = []
    for step in range(WINDOW_SIZE, len(data) - 1):
        current_prices = [data.iloc[step][f"close_{asset_type}"] for asset_type in ASSET_CODES.keys()]
        position_value = np.sum(positions * current_prices)
        total_asset = cash + position_value
        history.append({"total_asset": total_asset})

    return history


def main():
    # 1. 获取并预处理数据
    print("获取数据中...")
    data = get_multi_asset_data()
    print(f"数据获取完成，共{len(data)}个交易日")

    # 2. 初始化环境和智能体
    env = TradingEnv(data)
    state_dim = env._get_state().shape[0]  # 状态维度
    action_dim = 3 * env.n_assets  # 动作维度（3个动作×n个资产）
    agent = DQNAgent(state_dim, action_dim)

    # 3. 模型训练
    print("开始训练...")
    total_steps = 0
    train_loss = []
    while True:
        state = env.reset()
        done = False
        episode_loss = []
        while not done:
            # 选择动作
            action = agent.select_action(state)
            # 执行动作
            next_state, reward, done, info = env.step([action // env.n_assets] * env.n_assets)  # 简化：所有资产执行同一动作类型
            # 存储经验
            agent.store_transition(state, action, reward, next_state, done)
            # 学习
            loss = agent.learn()
            if loss is not None:
                episode_loss.append(loss)
            # 更新状态
            state = next_state
            total_steps += 1

            # 更新目标网络
            if total_steps % TARGET_UPDATE == 0:
                agent.update_target_net()

        if episode_loss:
            train_loss.append(np.mean(episode_loss))
        print(f"训练步数：{total_steps}，当前总资产：{info['total_asset']:.2f}")
        if total_steps > 30000:  # 训练终止条件：步数达到5000
            break

    # 4. 回测评估
    print("开始回测...")
    # DQN策略回测
    dqn_history = env.history
    dqn_metrics = calculate_metrics(dqn_history)
    # 买入并持有策略回测
    bah_history = backtest_buy_and_hold(data)
    bah_metrics = calculate_metrics(bah_history)

    # 5. 结果可视化
    plt.figure(figsize=(12, 6))
    # 资金曲线对比
    plt.subplot(1, 2, 1)
    dqn_assets = [h["total_asset"] for h in dqn_history]
    bah_assets = [h["total_asset"] for h in bah_history]
    plt.plot(dqn_assets, label="DQN策略")
    plt.plot(bah_assets, label="买入并持有")
    plt.xlabel("交易日")
    plt.ylabel("总资产（元）")
    plt.title("资金曲线对比")
    plt.legend()

    # 训练损失
    plt.subplot(1, 2, 2)
    plt.plot(train_loss)
    plt.xlabel("训练轮次")
    plt.ylabel("损失值")
    plt.title("DQN模型训练损失")
    plt.tight_layout()
    plt.show()

    # 6. 输出指标对比
    print("\n=== 回测指标对比 ===")
    print("DQN策略：", dqn_metrics)
    print("买入并持有：", bah_metrics)


if __name__ == "__main__":
    main()